#include <stdio.h>
#include <string.h>
#include <cv.h>
#include <cvaux.h>
#include <highgui.h>
#include <math.h>
#include <getopt.h>

// ścieżka do pliku detektora
#define DETECTOR "./haarcascade_frontalface_default.xml"

// ścieżka do pliku uczącego
#define TRAIN_FILE "./train.txt"

// ścieżka do pliku testowego
#define TEST_FILE "./test.txt"

// ścieżka do pliku z danymi twarzy własnych
#define EIGENFACES "./facedata.xml"

// liczba prób wykrycia twarzy przed zgłoszeniem błędu
#define ATTEMPTS 5

// czas podglądu
#define PREVIEW_TIME 3000

// domyślne ustawienia
int _DEVICE = CV_CAP_ANY;
int _PICTURES = 5; // liczba zdjęć
int _Y = 30; // zwiększenie wykrytego obszaru
int _SIZE_Y = 100; // rozmiar zdjęcia po skalowaniu
int _PREVIEW = 0; // podgląd

// zmienne pomocnicze
int nTrainFaces = 0;
int nEigens = 0;
IplImage** faceImgArr = 0;
CvMat* personNumTruthMat = 0;
IplImage* pAvgTrainImg = 0;
IplImage** eigenVectArr = 0;
CvMat* eigenValMat = 0;
CvMat* projectedTrainFaceMat = 0;

// deklaracje funkcji
void train();
void test();
void doPCA();
void storeTrainingData();
int loadTrainingData(CvMat** pTrainPersonNumMat);
int findNearestNeighbour(float* projectedTestFace, double* distance);
int* findNearestGroup(float* projectedTestFace, int groupSize);
int loadFaceImgArray(char* filename);
void query(int classID, int groupSize);
void printUsage();

IplImage* detect(IplImage* pInpImg, int preview);
IplImage* snapshot();
void collect(int id, int start);
void recognize();
void show(IplImage* pImage);

int main(int argc, char* argv[]) {
	int c;
	int classID, groupSize;
	char* action = NULL;

	opterr = 0;

	while((c = getopt(argc, argv, "a:d:e:s:c:n:p")) != -1) {
		switch(c) {
			case 'a':
				action = optarg;
				break;
			case 'c':
				classID = atoi(optarg);
				break;
			case 'd':
				_DEVICE = atoi(optarg);
				break;
			case 'e':
				_Y = atoi(optarg);
				break;
			case 'n':
				groupSize = atoi(optarg);
				break;
			case 's':
				_SIZE_Y = atoi(optarg);
				break;
			case 'p':
				_PREVIEW = 1;
				break;
			case '?':
				fprintf(stderr, "Invalid command or no argument specified. Type \"eigenfaces --help\" to view valid commands.\n", optopt);
				exit(1);
			default:
				abort();
		}
	}

	if(action == NULL) {
		printUsage();
		return 0;
	}
	
	if(strcmp(action, "train") == 0)
		train();
	else if(strcmp(action, "test") == 0)
		test();
	else if(strcmp(action, "recognize") == 0)
		recognize();
	else if(strcmp(action, "query") == 0)
		query(classID, groupSize);	
	else
		printUsage();
			
	return 0;
}

void printUsage() {
	fprintf(stderr, "Usage: eigenfaces <commands>\n"
			"Valid commands are\n"
			"\t-a\tAction: train, test, recognize, query. E.g. -a recognize\n"
			"\t-c\tPerson's class identifier. E.g. -c 2\n"
			"\t-d\tCamera device number. Value X of /dev/videoX . E.g. -d 0\n"
			"\t-e\tAdditional pixels for region of interest. E.g. -e 30\n"
			"\t-n\tFalse acceptance ratio. Finds \"n\" most similar pictures. E.g. -n 5\n"
			"\t-s\tSize of the image. E.g. -s 100\n"
			"\t-p\tDisplay preview. E.g. -p\n"
		);
}

void train() {
	int i;
	
	nTrainFaces = loadFaceImgArray(TRAIN_FILE);
	if(nTrainFaces < 2) {
		fprintf(stderr, "Need 2 or more training faces\n"
						"Input file contains only %d\n", nTrainFaces);
		return;
	}
	
	doPCA();
	
	projectedTrainFaceMat = cvCreateMat(nTrainFaces, nEigens, CV_32FC1);
	for(i = 0; i < nTrainFaces; ++i) {
		cvEigenDecomposite(faceImgArr[i], nEigens, eigenVectArr, 0, 0, pAvgTrainImg, projectedTrainFaceMat->data.fl + i*nEigens);
	}
	
	storeTrainingData();
}
		
int loadFaceImgArray(char* filename) {
	FILE* imgListFile = 0;
	char imgFilename[512];
	int iFace, nFaces = 0;
	
	imgListFile = fopen(filename, "r");
	
	while(fgets(imgFilename, 512, imgListFile))
		++nFaces;
	
	rewind(imgListFile);
	
	faceImgArr = (IplImage**)cvAlloc(nFaces*sizeof(IplImage*));
	personNumTruthMat = cvCreateMat(1, nFaces, CV_32SC1);
	
	for(iFace = 0; iFace < nFaces; ++iFace) {
		fscanf(imgListFile, "%d %s", personNumTruthMat->data.i + iFace, imgFilename);
		
		faceImgArr[iFace] = cvLoadImage(imgFilename, CV_LOAD_IMAGE_GRAYSCALE);
	}
	
	fclose(imgListFile);
	
	return nFaces;
}

void doPCA() {
	int i;
	CvTermCriteria calcLimit;
	CvSize faceImgSize;
	
	nEigens = nTrainFaces - 1;
	
	faceImgSize.width = faceImgArr[0]->width;
	faceImgSize.height = faceImgArr[0]->height;
	
	eigenVectArr = (IplImage**)cvAlloc(sizeof(IplImage*) * nEigens);
	
	for(i = 0; i < nEigens; ++i)
		eigenVectArr[i] = cvCreateImage(faceImgSize, IPL_DEPTH_32F, 1);
	
	eigenValMat = cvCreateMat(1, nEigens, CV_32FC1);
	
	pAvgTrainImg = cvCreateImage(faceImgSize, IPL_DEPTH_32F, 1);
		
	calcLimit = cvTermCriteria(CV_TERMCRIT_ITER, nEigens, 1);
	
	cvCalcEigenObjects(nTrainFaces, (void*)faceImgArr, (void*)eigenVectArr, CV_EIGOBJ_NO_CALLBACK, 0, 0, &calcLimit, pAvgTrainImg, eigenValMat->data.fl);
}

void storeTrainingData() {
	CvFileStorage* fileStorage;
	int i;
	
	fileStorage = cvOpenFileStorage(EIGENFACES, 0, CV_STORAGE_WRITE);
	
	cvWriteInt(fileStorage, "nEigens", nEigens);
	cvWriteInt(fileStorage, "nTrainFaces", nTrainFaces);
	cvWrite(fileStorage, "trainPersonNumMat", personNumTruthMat, cvAttrList(0, 0));
	cvWrite(fileStorage, "eigenValMat", eigenValMat, cvAttrList(0, 0));
	cvWrite(fileStorage, "projectedTrainFaceMat", projectedTrainFaceMat, cvAttrList(0, 0));
	cvWrite(fileStorage, "avgTrainImg", pAvgTrainImg, cvAttrList(0, 0));
	for(i = 0; i < nEigens; ++i) {
		char varname[200];
		sprintf(varname, "eigenVect_%d", i);
		cvWrite(fileStorage, varname, eigenVectArr[i], cvAttrList(0, 0));
	}
	
	cvReleaseFileStorage(&fileStorage);
}

void test() {
	int i, nTestFaces = 0;
	CvMat* trainPersonNumMat = 0;
	float* projectedTestFace = 0;
	
	nTestFaces = loadFaceImgArray(TEST_FILE);
	printf("%d test faces loaded\n", nTestFaces);
	
	if(!loadTrainingData(&trainPersonNumMat))
		return;
				
	projectedTestFace = (float*)cvAlloc(nEigens * sizeof(float));
	for(i = 0; i < nTestFaces; ++i) {
		int iNearest, nearest, truth;
		double distance;
		
		cvEigenDecomposite(faceImgArr[i], nEigens, eigenVectArr, 0, 0, pAvgTrainImg, projectedTestFace);
		
		iNearest = findNearestNeighbour(projectedTestFace, &distance);	
				
		truth = personNumTruthMat->data.i[i];
		nearest = trainPersonNumMat->data.i[iNearest];
		
		printf("%s\tnearest = %d, truth = %d, distance = %f\n", (nearest==truth?"OK":"FAIL"), nearest, truth, distance);
	}
}

void recognize() {
	int iNearest, nearest, photo = 0;	
	CvMat* trainPersonNumMat = 0;
	float* projectedTestFace = 0;
	double distance;
	IplImage* pImg = 0;
	IplImage* pFace = 0;
		
	do {
		pImg = snapshot();	
		if(!pImg)
			continue;
		pFace = detect(pImg, _PREVIEW);
		if(++photo >= ATTEMPTS) {
			printf("-1\n");
			fprintf(stderr, "Can't detect face");
			return;
		}
	} while(pFace == 0);
				
	if(!loadTrainingData(&trainPersonNumMat))
		return;
		
	projectedTestFace = (float*)cvAlloc(nEigens * sizeof(float));
			
	cvEigenDecomposite(pFace, nEigens, eigenVectArr, 0, 0, pAvgTrainImg, projectedTestFace);
		
	iNearest = findNearestNeighbour(projectedTestFace, &distance);
	nearest = trainPersonNumMat->data.i[iNearest];
		
	printf("%d\n", nearest);
	
	cvReleaseImage(&pImg);
	cvReleaseImage(&pFace);
}

void query(int classID, int groupSize) {
	int i, photo = 0;	
	int* iNearestGroup;
	CvMat* trainPersonNumMat = 0;
	float* projectedTestFace = 0;
	IplImage* pImg = 0;
	IplImage* pFace = 0;
		
	do {
		pImg = snapshot();	
		if(!pImg)
			continue;
		pFace = detect(pImg, _PREVIEW);
		if(++photo >= ATTEMPTS) {
			printf("-1\n");
			fprintf(stderr, "Can't detect face");
			return;
		}
	} while(pFace == 0);
				
	if(!loadTrainingData(&trainPersonNumMat))
		return;
		
	projectedTestFace = (float*)cvAlloc(nEigens * sizeof(float));
			
	cvEigenDecomposite(pFace, nEigens, eigenVectArr, 0, 0, pAvgTrainImg, projectedTestFace);
		
	iNearestGroup = findNearestGroup(projectedTestFace, groupSize);
	
	int found = 0;
	
	for(i = 0; i < groupSize; ++i) {
		if(trainPersonNumMat->data.i[iNearestGroup[i]] == classID) {
			found = 1;
			break;
		}
	}
	
	free(iNearestGroup);
		
	printf("%d\n", (found ? classID : -1));
	
	cvReleaseImage(&pImg);
	cvReleaseImage(&pFace);
}

int loadTrainingData(CvMat** pTrainPersonNumMat) {
	CvFileStorage* fileStorage;
	int i;
	
	fileStorage = cvOpenFileStorage(EIGENFACES, 0, CV_STORAGE_READ);
	if(!fileStorage) {
		printf("-1\n");
		fprintf(stderr, "Can't open "EIGENFACES);
		exit(1);
	}
	
	nEigens = cvReadIntByName(fileStorage, 0, "nEigens", 0);
	nTrainFaces = cvReadIntByName(fileStorage, 0, "nTrainFaces", 0);
	*pTrainPersonNumMat = (CvMat*)cvReadByName(fileStorage, 0, "trainPersonNumMat", 0);
	eigenValMat = (CvMat*)cvReadByName(fileStorage, 0, "eigenValMat", 0);
	projectedTrainFaceMat = (CvMat*)cvReadByName(fileStorage, 0, "projectedTrainFaceMat", 0);
	pAvgTrainImg = (IplImage*)cvReadByName(fileStorage, 0, "avgTrainImg", 0);
	eigenVectArr = (IplImage**)cvAlloc(nTrainFaces * sizeof(IplImage*));
	for(i = 0; i < nEigens; ++i) {
		char varname[200];
		sprintf(varname, "eigenVect_%d", i);
		eigenVectArr[i] = (IplImage*)cvReadByName(fileStorage, 0, varname, 0);
	}
	
	cvReleaseFileStorage(&fileStorage);
	
	return 1;
}

int findNearestNeighbour(float* projectedTestFace, double* distance) {
	*distance = DBL_MAX;
	int i, iTrain, iNearest = 0;
	
	for(iTrain = 0; iTrain < nTrainFaces; ++iTrain) {
		double distSq = 0;
		
		for(i = 0; i < nEigens; ++i) {
			float d_i = projectedTestFace[i] - projectedTrainFaceMat->data.fl[iTrain*nEigens + i];
			distSq += d_i * d_i;			
		}		
		
		if(distSq < *distance) {
			*distance = distSq;
			iNearest = iTrain;
		}	
	}
	
	return iNearest;
}

int* findNearestGroup(float* projectedTestFace, int groupSize) {
	int i, iTrain;
	int maxId;
	
	double* distTab = (double*)malloc(nTrainFaces*sizeof(double));
	int* groupTab = (int*)malloc(groupSize*sizeof(int));
	
	for(iTrain = 0; iTrain < nTrainFaces; ++iTrain) {
		double distSq = 0;
		
		for(i = 0; i < nEigens; ++i) {
			float d_i = projectedTestFace[i] - projectedTrainFaceMat->data.fl[iTrain*nEigens + i];
			distSq += d_i * d_i;			
		}		
		
		distTab[iTrain] = distSq;
	}
	
	for(i = 0; i < groupSize; ++i) {
		maxId = 0;
		for(iTrain = 1; iTrain < nTrainFaces; ++iTrain) {
			if(distTab[iTrain] < distTab[maxId]) {
				maxId = iTrain;
			}
		}	
		
		distTab[maxId] = DBL_MAX;
		groupTab[i] = maxId;
	}
	
	return groupTab;
}

IplImage* detect(IplImage* pInpImg, int preview) {
	CvHaarClassifierCascade * pCascade = 0; 
	CvMemStorage* pStorage = 0;
	CvSeq* pFaceRectSeq;
	IplImage* pFaceImg = 0;
	int i;
	
	pStorage = cvCreateMemStorage(0);
	pCascade = (CvHaarClassifierCascade*)cvLoad(DETECTOR, 0, 0, 0);
	
	if(!pInpImg || !pStorage || !pCascade)
		return 0;
	
	pFaceRectSeq = cvHaarDetectObjects(pInpImg, pCascade, pStorage, 1.1, 2, CV_HAAR_DO_CANNY_PRUNING, cvSize(60, 60));
		
	if(pFaceRectSeq->total > 0) {
		CvRect* pFaceRect = (CvRect*)cvGetSeqElem(pFaceRectSeq, 0);
				
		// powiększenie obszaru
		const int w = _Y;
		const int h = _Y;
		pFaceRect->x -= w;
		pFaceRect->y -= h;
		pFaceRect->width += 2*w;
		pFaceRect->height += 2*h;			
		// przesunięcie ROI w górę
		pFaceRect->y -= floor(0.09 * pFaceRect->height);			
		// zmniejszenie szerokości
		pFaceRect->width = 0.8*pFaceRect->width;
		pFaceRect->x += pFaceRect->height*0.1;		
		const int _SIZE_X = 0.8*_SIZE_Y;	
		
		if(preview) {
			IplImage* pPreview = cvCloneImage(pInpImg);
			CvPoint pt1 = { pFaceRect->x, pFaceRect->y };
			CvPoint pt2 = { pFaceRect->x + pFaceRect->width, pFaceRect->y + pFaceRect->height };
			cvRectangle(pPreview, pt1, pt2, CV_RGB(0,255,0), 3, 4, 0);
			show(pPreview);
		}
		
		cvSetImageROI(pInpImg, *pFaceRect);
		IplImage* pColorFaceImg = cvCreateImage(cvSize(_SIZE_X, _SIZE_Y), pInpImg->depth, pInpImg->nChannels);
		pFaceImg = cvCreateImage(cvGetSize(pColorFaceImg), IPL_DEPTH_8U, 1);
		
		cvResize(pInpImg, pColorFaceImg, CV_INTER_AREA);
		cvCvtColor(pColorFaceImg, pFaceImg, CV_BGR2GRAY);
		
		cvResetImageROI(pInpImg);
	}
	
	if(pCascade)
		cvReleaseHaarClassifierCascade(&pCascade);
	
	if(pStorage)
		cvReleaseMemStorage(&pStorage);
		
	return pFaceImg;
}

IplImage* snapshot() {
	int i;
	CvCapture* pCapture = 0;
	IplImage* pVideoFrame = 0;
	IplImage* pSnapShot = 0;
	
	pCapture = cvCaptureFromCAM(_DEVICE);
	if(!pCapture) {
		printf("-2\n");
		fprintf(stderr, "Failed to initialize video capture\n");
		exit(1);
	}
	
	for(i = 0; i < 25; ++i)
		pVideoFrame = cvQueryFrame(pCapture);

	pVideoFrame = cvQueryFrame(pCapture);
	if(!pVideoFrame) {
		printf("-2\n");
		fprintf(stderr, "Failed to get a video frame\n");
		exit(1);
	}

	pSnapShot = cvCloneImage(pVideoFrame);

	cvReleaseCapture(&pCapture);	
	
	return pSnapShot;
}
	
void show(IplImage* pImage) {
	cvNamedWindow("Image Preview", CV_WINDOW_AUTOSIZE);
	cvShowImage("Image Preview", pImage);
	cvWaitKey(PREVIEW_TIME);
	cvDestroyWindow("Image Preview");
}

